Enable Ascend NPU Backend with Custom Ops Integration for NF4 Support#1695
Enable Ascend NPU Backend with Custom Ops Integration for NF4 Support#1695SlightwindSec wants to merge 12 commits intobitsandbytes-foundation:mainfrom
Conversation
Signed-off-by: SlightwindSec <slightwindsec@gmail.com>
Signed-off-by: SlightwindSec <slightwindsec@gmail.com>
SlightwindSec
changed the title
|
Error Summary: Encountered a vector core execution failure on Ascend 910B3 NPU while running Qwen-image NF4 quantized model inference. The NPU reported multiple DDR memory access violations (error code 0x800000) across 12 compute cores, specifically during dequantize_blockwise_fp32_nf4_1kernel execution. The system threw ACL synchronization error (code 507035) when attempting tensor device transfer (pos_freqs.to(device)). Technical Breakdown: Hardware-Level: Multiple cores (5-15,20-22) triggered MTE (Memory Tagging Extension) faults indicating invalid DDR address access
|
…nabled Signed-off-by: SlightwindSec <slightwindsec@gmail.com>
Signed-off-by: SlightwindSec <slightwindsec@gmail.com>
|
@SlightwindSec Thanks for your fix. I made several test cases, and they all ran successfully without any bugs. Environment:
InferenceI tested both LLaMA3-8B and Qwen3-4B models. from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from bitsandbytes.functional import dequantize_4bit
import torch
quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4")
model_id = "./llama3_8b"
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="npu:0"
trust_remote_code=True,
dtype=torch.bfloat16,
quantization_config=quantization_config,
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
text = "Hello my name is"
inputs = tokenizer(text, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))TrainingFor training, I used the LLaMA-Factory framework to perform QLoRA fine-tuning with DDP and DeepSpeed ZeRO-3. You can reproduce the setup using the script below: |
SlightwindSec
changed the title
|
@SlightwindSec @matthewdouglas I have updated docs for install and README, would you mind having a check? |
|
@matthewdouglas Hi, the conflict in the CMakeLists.txt file have been resolved. Would you mind reviewing the changes and sharing any feedback or suggestions?. As mentioned in #1545, multi-Backend branch was deprecated, we are planning to merge ascend npu 4bit to main branch. |
TimDettmers
left a comment
TimDettmers
left a comment
There was a problem hiding this comment.
PR Review: #1695 — Enable Ascend NPU Backend with Custom Ops Integration for NF4 Support
Classification: [feature] [platform] [build] — New NPU backend with native AscendC kernels, build system changes, and device dispatch integration.
Size: 562 additions, 5 deletions across 12 files. Medium-large, new backend PR with new files.
This PR adds an Ascend NPU backend supporting NF4 4-bit quantization and dequantization. It includes a Python-side quantize_4bit implementation, native AscendC dequantization kernels (npu_kernels.cpp), CMake build integration, and cextension.py/__init__.py wiring. The scope is focused on NF4 only (no FP4, no 8-bit ops, no optimizers).
Blocking issues (4):
-
cextension.py: NPU detection takes priority over CUDA — In theBNB_BACKENDdetection chain,is_npu_available()is inserted beforetorch.cuda.is_available(). Sinceis_npu_available()only checks whether thetorch_npupackage is installed (viaimportlib.util.find_spec), NOT whether an actual NPU device is present, any system that hastorch_npuinstalled alongside CUDA will incorrectly setBNB_BACKEND = "NPU"instead of"CUDA". The XPU backend checkstorch._C._has_xpu(actual device presence); the HPU backend checkstorch.hpu.is_available(). The NPU backend needs an equivalent device-availability check, not just a package-installed check. Similarly, inget_native_library(), the NPU path would be selected over XPU when both packages are installed. -
cextension.py: Brokenif/elifchain inget_native_library()— The original code has two separateifblocks:if cuda_specs: ... binary_path = cuda_binary_pathfollowed by a separateif torch._C._has_xpu: .... The PR converts the second block's preceding blank line into anelif is_npu_available():clause, making the XPU check a separateifthat runs unconditionally. The result: on XPU systems withtorch_npuinstalled, the NPU path would be set, then immediately overwritten by the XPU path. But on NPU systems without XPU, the flow is correct only by coincidence. The library loading logic needs properif/elif/elifchaining across all backends. -
npu_ops.cpp: Memory leak —tilingHostis allocated withmalloc()but never freed. Every call todequantizeBlockwiseNf4()leakssizeof(BlockwiseNf4TilingData)bytes. This is called on every dequantization operation (every Linear4bit forward pass), so the leak accumulates during inference and training. -
npu_ops.cpp: Hardcoded core count —coreNumis hardcoded to 40 in both the tiling function and the kernel launchblockDim. This is specific to one Ascend chip variant (e.g., 910B). Different Ascend models have different core counts. This should either be queried at runtime or made configurable.
Additional issues (non-blocking but significant):
-
__init__.py: Redundant and insufficient availability check — The import guardif importlib.util.find_spec("torch") and importlib.util.find_spec("torch_npu")checks for package installation, not device availability. Thefind_spec("torch")check is redundant sincetorchis already imported at the top of the file. Other backends use device-availability checks (e.g.,torch.xpu.is_available(),torch.hpu.is_available()). The NPU backend should check for actual NPU device availability, e.g.,if importlib.util.find_spec("torch_npu"): import torch_npu; if torch.npu.is_available(): ...(following the HPU pattern). -
quantize_4bitis pure PyTorch, not a native kernel — The NPUquantize_4bitimplementation is a pure Python/PyTorch implementation, whiledequantize_4bituses native AscendC kernels. This asymmetry means quantization may be slow. The PR description acknowledges this ("temporary PyTorch version"), but it should be documented with a TODO in code. -
quantize_4bitmutates global state —global _NF4_QUANT_TABLE; _NF4_QUANT_TABLE = _NF4_QUANT_TABLE.to(A.device)modifies the shared module-level_NF4_QUANT_TABLEfrombackends/utils.py. This is not thread-safe and would break if multiple devices are used simultaneously. Other backends using this table should be checked for the same pattern. -
NF4 packing order differs from CUDA — The comment in
quantize_4bitstates "Pack 4-bit values in NPU-compatible order (low nibble first) to match NPU-specific unpacking logic; differs from CUDA's packing." This means NPU-quantized checkpoints are NOT compatible with CUDA-dequantization and vice versa. This is a significant interoperability concern: users who quantize on NPU cannot dequantize on CUDA. This should be prominently documented. -
No tests included — The PR adds no test files. For a new backend, there should be at minimum integration tests for
quantize_4bitanddequantize_4bitround-trip correctness, and parametrized device tests. -
npu_kernels.cpp: Missing newline at end of file. -
cextension.py:import importlibplacement — The import is added after thefrom typing import Optionalline but before theimport torchline, violating the standard library grouping convention (it should be grouped with the other stdlib imports at the top, not aftertyping). -
CMake:
npu-smiparsing is fragile — The CMake auto-detection runsnpu-smi info | awk ... 'NR==7 {print $3}'which relies on exact output format ofnpu-smi. Different CANN versions or locales may change the output format. The error message also has a typo: "Auto-detech" should be "Auto-detect".
Security assessment:
- No network access, no command execution, no serialization exploits, no obfuscation patterns in the Python code.
- The native C++ code (
npu_ops.cpp,npu_kernels.cpp) uses AscendC SDK APIs (ACLRT_LAUNCH_KERNEL,kernel_operator.h). These are vendor-provided APIs analogous to CUDA runtime APIs. The code does not execute arbitrary commands or access the filesystem. - The CMake changes execute
npu-smiduring build configuration (analogous torocminfofor ROCm). This is a build-time tool invocation, consistent with existing patterns. pythonInterface.cppchanges are properly guarded behind#if BUILD_NPUpreprocessor conditionals.- The
import importlibaddition is benign (already used elsewhere in the codebase). - No security concerns identified.
Downstream Impact
Risk level: LOW (but see blocking issue #1)
The PR adds a new backend without changing existing APIs. The matmul_4bit change in _functions.py adds and A.device.type != "npu" to an existing condition, which is additive and follows the same pattern used for HPU.
However, blocking issue #1 (NPU taking priority over CUDA when torch_npu is installed) could silently break CUDA users who happen to have torch_npu installed. This would be a regression for existing users.
- Transformers: Not affected (additive backend)
- PEFT: Not affected
- Accelerate: Not affected
- TGI: Not affected
- vLLM: Not affected
Cross-PR Conflicts
- PR #1853 (Add MPS backend): overlaps on
__init__.py,README.md— both add new backend imports and README table entries. File-level conflicts likely. - PR #1826 (Metal backend with libtorch): overlaps on
cextension.py,__init__.py,CMakeLists.txt— similar new-backend additions. - PR #1823 (Add mps backend python only): overlaps on
cextension.py,__init__.py. - PR #1858 (Add k-bit blockwise quantization): overlaps on
CMakeLists.txt,csrc/pythonInterface.cpp.
None of these are semantic conflicts — they are independent backend/feature additions that will need mechanical merge resolution.
Verdict: Request changes. The CUDA-priority regression (issue #1), broken if/elif chain (issue #2), memory leak (issue #3), and hardcoded core count (issue #4) must be fixed before this can merge. The missing tests (issue #9) are also a blocking concern for a new backend.
- Security: Clear
- Downstream impact: LOW (additive backend), but issue #1 risks CUDA regression
- Tests: Missing — no tests included for the new backend
- CI: Not triggered (fork PR, needs maintainer approval)
- Serialization: NPU packing format differs from CUDA (cross-device checkpoint incompatibility)
- Cross-PR conflicts: File-level conflicts with 4 open PRs (mechanical resolution)
| if torch.version.hip: | ||
| HIP_ENVIRONMENT = True | ||
| BNB_BACKEND = "ROCm" | ||
| elif is_npu_available(): |
There was a problem hiding this comment.
Blocking: NPU detection takes priority over CUDA. is_npu_available() only checks if the torch_npu package is installed (via importlib.util.find_spec), not whether an actual NPU device is present. Since this elif comes before torch.cuda.is_available(), any system with torch_npu installed alongside CUDA will incorrectly set BNB_BACKEND = "NPU". This would silently break all CUDA users who happen to have torch_npu in their environment.
The fix: check for actual device availability (e.g., torch.npu.is_available() after importing torch_npu), and place the NPU check after CUDA in the priority chain, following the pattern of XPU/HPU.
| binary_path = cuda_binary_path | ||
|
|
||
| elif is_npu_available(): | ||
| binary_path = PACKAGE_DIR / f"libbitsandbytes_npu{DYNAMIC_LIBRARY_SUFFIX}" |
There was a problem hiding this comment.
Blocking: Broken if/elif chain. The original code had two independent if blocks here (one for cuda_specs, one for _has_xpu). By inserting elif is_npu_available() between them, the XPU check is now disconnected from the chain. The result is that on a system with both torch_npu and XPU, the NPU binary path is set, then immediately overwritten by the XPU path in the next if. This needs proper if/elif/elif/elif chaining across all backends.
| BlockwiseNf4TilingData *tilingHost; | ||
| tilingHost = (struct BlockwiseNf4TilingData *)malloc(tilingSize); | ||
| uint32_t error = get_dequantize_blockwise_nf4_tiling(blocksize, n, tilingHost); | ||
| if (error != 0) { |
There was a problem hiding this comment.
Blocking: Memory leak. tilingHost is allocated with malloc() but never freed. This leaks memory on every dequantization call (every Linear4bit.forward()). Add free(tilingHost) after the kernel launch, or use stack allocation since BlockwiseNf4TilingData is a small struct.
| tiling->ubSize = 196 * 1024; | ||
| uint32_t coreNum = 40; | ||
| uint32_t totalPkgNum = (n + blocksize - 1) / blocksize; | ||
| uint32_t singleCorePkgNum = (totalPkgNum + coreNum - 1) / coreNum; |
There was a problem hiding this comment.
Blocking: Hardcoded core count. coreNum = 40 is specific to one Ascend chip variant. Different Ascend models have different AI core counts (e.g., Ascend 310 has 2 cores, 910A has 32, 910B has 40). This should be queried from the hardware at runtime or passed as a parameter.
5 participants
What does this PR do?
This PR ports Ascend NPU backend changes from the multi-backend-refactor branch and integrates with custom ops. It includes changes to enable Ascend build and translation of kernels and ops to Ascend-compatible operators. As the AscendC-based high-performance NF4 implementation is still in progress, a temporary PyTorch version is used for now. The build steps remain the same as before from the user's standpoint.
Collaborators
@ji-huazhong @Ginray @Runningwater23 @jiaqiw09
cc @Titus-von-Koeller @matthewdouglas @amathews-amd @sunway513